from tt.scenarios import *
from collect.trace_collector import collect_traces_all
from locust import HttpUser, task, events, between
import gevent
from locust.env import Environment
from locust.stats import stats_printer, stats_history
from locust.log import setup_logging
import time
from datetime import datetime

# setup_logging("INFO", None)

# class MyUser(HttpUser):
#     wait_time = between(1, 5)
#
#     @task(1)
#     def hello1(self):
#         self.client.get("/hello")
#
#     @task(2)
#     def world(self):
#         self.client.get("/world")
#
#     def on_start(self):
#         self.client.post("/login", json={"name": "dyz", "age": 18})

hostIP = "139.198.164.54"

class TTUser(HttpUser):

    host = "http://" + hostIP + ":32677"

    wait_time =  between(1, 5)

    def on_start(self):
        self.url = "http://139.198.164.54:32677"
        self.q = Query(self.url)
        if not self.q.login():
            logging.fatal('login failed')

    @task(1)
    def task_query_and_cancel(self):
        query_and_cancel(self.q)

    @task(1)
    def task_query_and_collect(self):
        query_and_collect(self.q)

    @task(1)
    def task_query_and_execute(self):
        query_and_execute(self.q)

    @task(5)
    def task_query_and_preserve(self):
        query_and_preserve(self.q)

    # @task(1)
    # def preserve_and_consign(self):
    #     query_and_consign(self.q)

    @task(2)
    def task_query_and_pay(self):
        query_and_pay(self.q)

    @task(3)
    def task_query_and_rebook(self):
        query_and_rebook(self.q)


# if __name__ == "__main__":
#
#     workload_duration = 600
#     num_user = 1
#     stop_duration = 60
#     host = "139.198.183.68"
#
#     timestamp_start = time.time()
#     time_start = datetime.fromtimestamp(timestamp_start).strftime('%Y-%m-%d %H%M%S')
#     print("载荷发生开始时间：{}".format(time_start))
#     print("载荷持续时间：" + str(workload_duration) + "秒，模拟用户数：" + str(num_user))
#
#     env = Environment(user_classes=[TTUser])
#     runner = env.create_local_runner()
#     web_ui = env.create_web_ui("127.0.0.1", 8089)
#     # execute init event handlers (only really needed if you have registered any)
#     env.events.init.fire(environment=env, runner=runner, web_ui=web_ui)
#
#     # start a greenlet that periodically outputs the current stats
#     gevent.spawn(stats_printer(env.stats))
#
#     # start a greenlet that save current stats to history
#     gevent.spawn(stats_history, env.runner)
#
#     # start the test
#     runner.start(num_user, spawn_rate=10)
#
#     # in 600 seconds stop the runner
#     gevent.spawn_later(workload_duration, lambda: runner.quit())
#
#     # wait for the greenlets
#     runner.greenlet.join()
#
#     # stop the web server for good measures
#     web_ui.stop()
#
#     timestamp_end = time.time()
#
#     time_end = datetime.fromtimestamp(timestamp_end).strftime('%Y-%m-%d %H%M%S')
#     print("载荷发生结束！时间为：{}".format(time_end))
#     print("等待{}秒后开始拉取trace数据".format(stop_duration), end="")
#
#     for i in range(stop_duration):
#         time.sleep(1)
#         print(".", end="")
#
#     collect_traces_all(hostIP, timestamp_start, timestamp_end)